package com.ai.model.tyqw;

import com.ai.model.ChatRequest;
import com.alibaba.dashscope.aigc.generation.Generation;
import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.dashscope.aigc.generation.models.QwenParam;
import com.alibaba.dashscope.common.Message;
import com.alibaba.dashscope.common.Role;
import io.reactivex.Flowable;
import org.springframework.stereotype.Component;

import java.util.Arrays;

/**
 * 通义千问大模型调用
 *
 * @author hh
 * @since 2024/6/22
 */
@Component
public class TyqwChatAiRequest implements ChatRequest {

    private final TyqwApiProperties tyqwApiProperties;

    public TyqwChatAiRequest(TyqwApiProperties tyqwApiProperties) {
        this.tyqwApiProperties = tyqwApiProperties;
    }

    @Override
    public String chat(String question) throws Exception {
        Generation gen = new Generation();
        Message userMsg = Message.builder().role(Role.USER.getValue()).content(question).build();
        QwenParam param = QwenParam.builder().model(tyqwApiProperties.getModel())
                .messages(Arrays.asList(userMsg))
                .resultFormat(QwenParam.ResultFormat.MESSAGE)
                .apiKey(tyqwApiProperties.getApiKey())
                .workspace(tyqwApiProperties.getWorkspaceId())
                .topP(0.8)
                .incrementalOutput(true) // get streaming output incrementally
                .build();
        Flowable<GenerationResult> result = gen.streamCall(param);
        StringBuilder fullContent = new StringBuilder();
        result.blockingForEach(item -> fullContent.append(item.getOutput().getChoices().get(0).getMessage().getContent()));
        return fullContent.toString();
    }

}
